from keras.models import Sequential
from keras.layers import *
from keras.callbacks import ModelCheckpoint
from keras.losses import MeanSquaredError
from keras.metrics import RootMeanSquaredError
from keras.optimizers import Adam
from keras.models import *
from keras.layers import *
from keras.layers import concatenate
import tensorflow as tf
from tcn import TCN


def lstm(x_train, y_train, x_val, y_val, x_test, y_test, windows, featureSize):
    # LSTM
    model_lstm = Sequential()
    model_lstm.add(InputLayer((7, 9)))
    # model_lstm.add(Conv1D(64, kernel_size=2))
    model_lstm.add(
        LSTM(64, activation='relu', input_shape=(x_train.shape[1], x_train.shape[2]), return_sequences=True))
    model_lstm.add(LSTM(32))
    # model_lstm.add(Dropout(0.01))
    # model_lstm.add(Attention(use_scale=False))
    # model_lstm.add(Dropout(0.1))
    model_lstm.add(Dense(8, 'relu'))
    model_lstm.add(Dense(1, 'linear'))

    # model_Multivariate.build(input_shape=(None, 64, 8, 1))
    # model_Multivariate.summary()

    # cp_lstm = ModelCheckpoint('model_lstm/', save_best_only=True)
    model_lstm.compile(loss=MeanSquaredError(), optimizer=Adam(learning_rate=0.0001), metrics=[RootMeanSquaredError()])
    history = model_lstm.fit(x_train, y_train, batch_size=32, validation_data=(x_val, y_val), epochs=120)